import os

import pandas as pd
from typing import *
import re

from openpyxl.utils import get_column_letter

from utils.tools import SingleTon
from openpyxl import Workbook
from openpyxl.styles import PatternFill
from utils.error import DataHandlerError


class Handler(metaclass=SingleTon):
    def __init__(self, uuid: str):
        self.uuid = uuid
        self.data = pd.DataFrame()

    def read(self, data: pd.DataFrame):
        pass

    def dump(self, path: str):
        pass


class PickHandler(Handler):
    """
    中性专用回测数据处理
    """
    factor_prefix: str = '_factor_'
    final_value_name: str = 'v'

    def __init__(self, uuid: str):
        super().__init__(uuid)
        self.factors: Dict[str, str] = {}
        self._extract_factors()

    def _extract_factors(self):
        for col in self.data.columns:
            res = re.match(rf'^{self.factor_prefix}(.*)$', col)
            if res:
                self.factors[res.groups()[0]] = col

    def read(self, data: pd.DataFrame):
        self.data = pd.concat([self.data, data])

    def dump(self, path: str):
        if self.data.empty:
            return

        self.data['open_time'] = pd.to_datetime(self.data['open_time'], unit='ms')
        select_coin = self.data[self.data['side'] != 0]
        df = self.data.loc[:, ['open_time', 'symbol', 'factor_rank', 'filter_rank']]
        # 使用set_index设置open_time和symbol为索引，然后将factor_rank和filter_rank unstack
        df_reshaped = df.set_index(['open_time', 'symbol']).stack().unstack(level=[0, 2])

        # 现在的列索引是 (open_time, variable)，其中 variable 是 factor_rank 或 filter_rank
        # 为了使列索引更清晰，我们可以重命名它
        df_reshaped.columns = pd.MultiIndex.from_tuples(df_reshaped.columns, names=['open_time', 'Rank Type'])
        # df_reshaped = df_reshaped.reset_index()
        begin_time_list = df['open_time'].astype(str).unique().tolist()
        symbol_list = sorted(df['symbol'].unique().tolist())
        with pd.ExcelWriter(os.path.join(path, f'{self.uuid}.pick.xlsx'), engine='openpyxl') as writer:
            # 将 DataFrame 保存到 Excel 文件中
            df_reshaped.to_excel(writer, sheet_name='Sheet1', index=True)
            worksheet = writer.book['Sheet1']
            # 定义一个颜色映射
            red_fill = PatternFill(start_color="FFFF0000", end_color="FFFF0000", fill_type="solid")
            green_fill = PatternFill(start_color="FF00FF00", end_color="FF00FF00", fill_type="solid")
            # count表示列坐标，excel从第二列开始表示时间
            count = 0
            # 统一调整列宽
            column_len = len(str(begin_time_list[0])) * 1.5
            for column in range(1, len(begin_time_list) + 1):
                column_letter = get_column_letter(column)
                worksheet.column_dimensions[column_letter].width = column_len
            for open_time in begin_time_list:
                count += 2
                original = select_coin[select_coin['open_time'] == open_time]
                original_symbol = original['symbol'].tolist()
                for symbol in original_symbol:
                    # 获取币种在表格元素中的index，因为list.index是从0开始，而excel从4开始，则+4
                    index = symbol_list.index(symbol)
                    cell = worksheet.cell(row=index + 4, column=count)
                    direction = original[original['symbol'] == symbol]['side'].values[0]
                    # 正向绿色，反向红色
                    if direction == 1:
                        cell.fill = green_fill
                    elif direction == -1:
                        cell.fill = red_fill


handler_map: Dict[str, type] = {
    'pick': PickHandler
}
